#!/usr/bin/env python
import sys,os,re
import ssummo.CONFIG as CONFIG
import dict_to_phyloxml
import ssummo.ssummolib as slib
import ssummo.colours as colours

class LocalOptions( slib.Options ):
    def __init__(self,args=None):
        self.options = { '-out'       : None,
                         '-in'        : '',
                         '-itol-type' : 'multivaluebar',
                         '-name-col'  : 0,
                         '-data-col'  : ['all'],
                         '--headers'  : True,
                         '-headers'   : [],
                         '-delimiter' : '\t',
                         }
        self.singleargs = ['-in','-out','-graph-type','-data-col','-name-col',]
        self.multiargs = ['-headers','-data-col']
        self.helpTxt = { '-itol-type' : 'IToL figure type. Currently only multivaluebar supported',
                         '-in'        : '',
                         '-out'       : 'Output file for prefix for phyloxml file and IToL data file',
                         '-name-col'  : 'Column index which contains the name (first col is 0).',
                         '-data-col'  : 'Space separated list of column numbers which contain data (first is 0).',
                         '--headers'  : 'If the file contains headers, use them in the field descriptions [True]',
                         '-headers'   : 'If the file does not contain headers, supply header info here',
                         '-delimiter' : 'Data delimiter',
                }
        self.post_checks += [self.integers,self.reqs]
        if args is not None:
            self.parseArgs(args)
    def integers(self):
        self['-name-col'] = int(self['-name-col'])
        if 'all' not in self['-data-col']:
            for col,entry in enumerate( self['-data-col'] ):
                self['-data-col'][col] = int(entry)
    def reqs(self):
        if self['-out'] is None:
             raise IOError( 'Need to define -out!\n' )

def separate( options ):
    file_name = options['-in']
    with file( file_name ,'r' ) as in_file:
        for line in in_file:
            row = line.rstrip().split( options['-delimiter'] )
            yield row

def join_results( results, taxpath, species ):
    path = taxpath.split( os.path.sep )
    if len( path ) >= 2:
        path = path[:-1]
    node = results
    for tax in path:
        if tax in node:
            node = node[tax]
        else:
            node.update( { tax : {} } )
            node = node[tax]
    node.update( { species : {} } )
    return results

def main(options):
    TDB = slib.TaxDB()
    tdict = slib.load_index()
    tdict_node = tdict
    splitter = separate( options )
    spaceSub = re.compile( r'\s+' )
    text_formatter = re.compile( '["\',/]' )
    with file( options['-out'] + '.txt' ,'w' ) as out_file:
        headers = splitter.next()
        if options['-data-col'] == ['all']:
            options['-data-col'] = range( 0 , options['-name-col'] ) + range( 1+options['-name-col'] , len( headers ) )
        get_data = lambda data: [ data[index] for index in options['-data-col'] ]
        print get_data( headers )
        out_file.write( 'LABELS,{0}\n'.format( ','.join(get_data(headers))) )
        cols = colours.generate_HEX_colours( len( options['-data-col'] ) )
        out_file.write( 'COLORS,{0}\n'.format( ','.join( cols ) ) )
        results = {}
        for line in splitter:
            species = text_formatter.sub( '' ,line[ options['-name-col'] ]) 
            # MAKE DICTIONARY
            print 'Searching for ',species,
            rank,name = slib.fetchRank( TDB, 'Prokaryotes' , species )
            print '. Found ',name,' at rank: ',rank
            if rank is not False:
                taxpath = slib.find_node_name( name , tdict_node )
                results = join_results( results, taxpath , species )

                data = get_data( line )
                out_file.write('{0},{1}\n'.format( spaceSub.sub('_',species) , ','.join(data) ))
            #data = []
            #for col in options['-data-col']:
            #    data.append( line[col] )
    with file( options['-out'] + '.xml','w' ) as xmlfile:
        print 'Saving phyloxml file to ',xmlfile.name
        dict_to_phyloxml.write_xml( CONFIG.top , xmlfile , results, tdict, )

if __name__ == '__main__':
    args = sys.argv[1:]
    options = LocalOptions(args)
    main( options )
